import csv
import collections
import StringIO

class DataTable:
    """A spreadsheet-ish table of data.
    
    Initialized with a csv string, whose first line provides labels
    for the columns.
    
    >>> north = DataTable("state,capital\\nMA,Boston")
    >>> south = DataTable("state,capital\\nGA,Atlanta")
    >>> (north + south).csv
    'state,capital\\nMA,Boston\\nGA,Atlanta\\n'
    """

    @property
    def csv(self):
        return self._csv

    def __init__(self,csv):
        self._csv = csv

    def __add__(self,other):
        def header_body(csv):
            file = StringIO.StringIO(csv)
            header = file.readline()
            body = file.read()
            if body[-1] != '\n':
                body += '\n'
            return (header,body)
        (self_header,self_body) = header_body(self.csv)
        (other_header,other_body) = header_body(other.csv)
        assert self_header == other_header
        # TODO: if it'd be useful, we could handle variation in column order.
        return DataTable(self_header + self_body + other_body)

    def summarize(self,group_cols,summary_cols):
        """Given suitable functions, returns a summary of the table.

        >>> n = DataTable("region,state,capital\\nNew England,MA,Boston")
        >>> n += DataTable("region,state,capital\\nNew England,RI,Providence")
        >>> mw = DataTable("region,state,capital\\nMidwest,MN,St. Paul\\nMidwest,WI,Madison")
        >>> (n + mw).csv
        'region,state,capital\\nNew England,MA,Boston\\nNew England,RI,Providence\\nMidwest,MN,St. Paul\\nMidwest,WI,Madison\\n'
        >>> group_cols = ['region']
        >>> summary_cols = { \
            'COUNT': \
                    lambda summary_row, data_row: \
                        int(summary_row['COUNT']) + 1, \
            'CONCAT': \
                    lambda summary_row, data_row: \
                        (summary_row['CONCAT'] or '')+ ':' +data_row['state'] \
            }
        >>> (n + mw).summarize(group_cols,summary_cols).csv
        'region,CONCAT,COUNT\\nNew England,:MA:RI,2\\nMidwest,:MN:WI,2'

        The arguments to summarize require some explanation:

        group_cols is simply a list of fields to group by in the summary.

        summary_cols is a dict relating summary column names to functions defining their values.
        The functions in turn take two arguments:
            summary_row is a dict whose values are used as accumulators.
            data_row is set to each row of the table in turn.

        TODO: it's kind of inconsistent about when it returns a string, when a number.
        TODO: would you ever refer to more than one key from summary row?
        """
        summary = collections.defaultdict(
            lambda: collections.defaultdict(int)
            # summary is a dict,
            #   every value of which is a dict
            #     where unspecified values default to zero.
        )
        csv_file = StringIO.StringIO(self.csv)
        for data_row in csv.DictReader(csv_file):
            key = tuple(
                map(
                    lambda col:
                        data_row[col],
                    group_cols
                )
            )
            value = dict(
                map(
                    lambda summary_col:
                        (summary_col,summary_cols[summary_col](summary[key],data_row)),
                    summary_cols.keys()
                )
            )
            summary[key]=value
        sorted_summary_cols_keys = sorted(summary_cols.keys())
        # TODO: use a library that would quote/escape commas in data
        header = ','.join(group_cols+sorted_summary_cols_keys)+'\n'
        rows = map(
            lambda key:
                ','.join(
                    list(key) + \
                    map(
                        lambda summary_col_key:
                            str(summary[key][summary_col_key]),
                        sorted_summary_cols_keys
                    )
                ),
            summary.keys()
        )
        csv_string = header+'\n'.join(rows)
        return DataTable(csv_string)


def from_csv_file(csv_filename):
    """Given the name of a csv file, returns a DataTable."""
    return DataTable(
        open(csv_filename,'r').read()
    )

if __name__ == "__main__":
    import doctest
    doctest.testmod()

